# -*- coding: utf-8 -*-
"""
    Time    : 2021/1/4 9:42 下午
    Author  : Thinkgamer
    File    : 698-划分为k个相等的子集.py
    Desc    : https://leetcode-cn.com/problems/partition-to-k-equal-sum-subsets/
"""
"""
给定一个整数数组 nums 和一个正整数 k，找出是否有可能把这个数组分成 k 个非空子集，其总和都相等。

示例 1：
输入： nums = [4, 3, 2, 3, 5, 2, 1], k = 4
输出： True
说明： 有可能将其分成 4 个子集（5），（1,4），（2,3），（2,3）等于总和。

提示：
1 <= k <= len(nums) <= 16
0 < nums[i] < 10000
"""


# 错误解法，因为覆盖不到所有的情况
def can_partition_k_subsets(nums, k):
	num_count_dict = dict()
	count = 0
	maxv = 0
	minv = float("inf")
	for num in nums:
		count += num
		if num > maxv:
			maxv = num
		if num < minv:
			minv = num
		num_count_dict.setdefault(num, 0)
		num_count_dict[num] += 1
	avg = count // k
	# 如果不能被整除 或 最大值 比 均值还大，说明不可划分为k个集合
	if count % k != 0 or maxv > avg:
		return False
	# 限制每个数字的使用次数
	num_use_dict = dict()
	for key, value in num_count_dict.items():
		num_use_dict[key] = value // k if value // k > 0 else 1
	result = list()
	print("avg is: {}, minv is: {}".format(avg, minv))
	print("num_count_dict: {}, num_use_dict: {}".format(num_count_dict, num_use_dict))
	for one in nums:
		print("one: {}, {}".format(one, num_use_dict))
		# 如果该值的剩余次数 小于 等于 0。则该值不进行遍历，因为在后续操作中可能会是 one 的剩余次数 减少
		if num_count_dict[one] <= 0:
			continue
		# 如果该值刚好等于 avg
		if one == avg:
			result.append((one))
			num_count_dict[one] -= 1
		# 如果不等于avg
		else:
			# 先假设 one 被使用上，进行减1操作
			num_count_dict[one] -= 1
			# one 在单词循环中的可使用次数 -1
			num_use_dict[one] -= 1
			diff = avg - one
			# 如果差值刚好存在
			if nums.__contains__(diff) and num_count_dict[diff] > 0 and num_use_dict[diff] > 0:
				result.append((one, diff))
				num_count_dict[diff] -= 1
				num_use_dict[diff] += 1
				num_use_dict[one] += 1
			# 如果差值不存在, 从 diff 开始依次缩小遍历
			else:
				diff -= 1
				temp = [one]
				use_diff = [one]
				while diff > 0:
					# 如果 diff - 1在 num_count_dict 中
					if nums.__contains__(diff) and num_count_dict[diff] > 0 and num_use_dict[diff] > 0:
						temp.append(diff)
						num_count_dict[diff] -= 1
						num_use_dict[diff] -= 1
						use_diff.append(diff)

						one = one + diff
						diff = avg - one
					elif diff > minv:
						diff -= 1
					else:
						break
				# print("diff: {}, num_count_dict: {}".format(diff, num_count_dict))
			
				if diff != 0:
					return False
				else:
					result.append(temp)
					for a in use_diff:
						num_use_dict[a] += 1
			
		print(result)
	for k,v in num_count_dict.items():
		if v != 0:
			return False
	return True
	

# 参考题解
# 52 ms 77.45% ｜ 15mb 14.49%
def can_partition_k_subsets_v2(nums, k):
	# target 表示 均值，rem 表示余数
	target, rem = divmod(sum(nums), k)
	if rem != 0:
		return False
	
	# 递归调用，判断是否可以用当前的nums数组去凑到k个target
	def search(groups):
		# 整个 nums 数组已经全部用完了，还没有从递归中 return 上来，说明可以完成这个任务
		if not nums:
			return True
		
		# 从 nums 中取出最大的数
		v = nums.pop()
		# 判断放到哪个组里
		for i, group in enumerate(groups):
			# 如果这个组加上这个数 v，还没有超过 target
			if group + v <= target:
				# 那就把这个数加到这个组里
				groups[i] += v
				# 继续判断后续能否完成任务，如果可以，直接返回True，这会在递归中一路返回上去
				# 在这里递归找到 v 应该放到哪个组里边
				if search(groups):
					return True
				# 否则，这个数不应该放到这个组里，把v从这个组中拿出来，继续尝试其他的组
				groups[i] -= v
			# 如果这个组有数字了，先不要往这个组放了，去放到下一个组
			# 感觉也有点启发式算法的意思，为什么要这么搞呢？因为数组已经被排序了，单调性原理
			if not group:
				break
		# 找了一圈发现 v没有找到对应的位置，必须把v放回去原来的数组
		nums.append(v)
		# 并且返回False 表示这次回溯是失败的，这个v没有找到对应位置
		return False
		
	# 对 nums 进行排序，方便进行快速的遍历
	nums.sort()
	if nums[-1] > target:
		return False
	# 如果有和目标值相等数，那么自己单独成为一组，同时在nums中将其移除
	while nums and nums[-1] == target:
		nums.pop()
		k -= 1
	return search([0] * k)

for nums, k in [
	# ([4, 3, 2, 3, 5, 2, 1], 4),
	# ([6, 1, 1, 1], 3),
	# ([1, 2, 3, 3, 4, 5], 3),
	# ([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 5),
	([10, 10, 10, 7, 7, 7, 7, 7, 7, 6, 6, 6], 3),
]:
	result = can_partition_k_subsets_v2(nums, k)
	print(result)
